#!/usr/bin/env python3
"""
批量 yt-dlp 下载脚本（支持 cookies-from-browser / cookiefile、断点续传、重试）
用法示例:
  python batch_ytdlp_downloader.py -i youtube_urls.txt --cookies-from-browser chrome
  python batch_ytdlp_downloader.py -i youtube_urls.txt --cookiefile /path/to/cookies.txt
"""

import argparse
import os
import sys
from yt_dlp import YoutubeDL

def build_opts(output_dir, cookiefile=None, cookies_from_browser=None, audio_only=False, max_height=None):
    outtmpl = os.path.join(output_dir, '%(title)s [%(id)s].%(ext)s')
    fmt = 'bestaudio/best' if audio_only else 'bestvideo+bestaudio/best'
    if max_height:
        fmt = f'bestvideo[height<={max_height}]+bestaudio/best[height<={max_height}]'
    opts = {
        'format': fmt,
        'outtmpl': outtmpl,
        'merge_output_format': 'mp4',
        'noplaylist': False,
        'ignoreerrors': True,
        'continuedl': True,
        'retries': 10,
        'fragment_retries': 10,
        'concurrent_fragment_downloads': 5,
        'http_chunk_size': 1048576,
        'download_archive': os.path.join(output_dir, 'downloaded.txt'),  # 跳过已下载的视频
        'no_warnings': False,
        'progress_hooks': [lambda d: print(f"[{d.get('status')}] {d.get('filename','')}")],
        # 避免被封杀：模拟常见 UA（可选）
        'http_headers': {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'},
    }
    if cookiefile:
        opts['cookiefile'] = cookiefile
    elif cookies_from_browser:
        # yt-dlp Python API：使用 'cookiesfrombrowser' 需要通过 'cookiesfrombrowser' key
        # 注意：有时直接在 Python API 不生效，CLI 用法更稳妥；但我们仍尝试提供。
        opts['cookiesfrombrowser'] = cookies_from_browser

    return opts

def read_urls(input_file):
    with open(input_file, 'r', encoding='utf-8') as f:
        lines = [l.strip() for l in f if l.strip() and not l.strip().startswith('#')]
    return lines

def main():
    p = argparse.ArgumentParser()
    p.add_argument('-i','--input', required=True, help='每行一个 URL 的文件（UTF-8）')
    p.add_argument('-o','--output', default='downloads', help='保存目录')
    p.add_argument('--cookiefile', default=None, help='Netscape 格式 cookies 文件路径（可选）')
    p.add_argument('--cookies-from-browser', default=None, help='直接从浏览器读取 cookies（chrome, firefox, edge, chromium）')
    p.add_argument('--audio-only', action='store_true', help='只下载音频')
    p.add_argument('--max-height', type=int, default=None, help='限制视频高度（例如 720）')
    args = p.parse_args()

    os.makedirs(args.output, exist_ok=True)
    urls = read_urls(args.input)
    if not urls:
        print("输入文件没有 URL。")
        sys.exit(1)

    ydl_opts = build_opts(args.output, cookiefile=args.cookiefile, cookies_from_browser=args.cookies_from_browser,
                         audio_only=args.audio_only, max_height=args.max_height)

    # CLI 通常在处理 cookies_from_browser 更稳定；如果 Python API 读取失败，提示用户用 CLI。
    print(f"准备下载 {len(urls)} 个链接 -> {args.output}")
    if args.cookies_from_browser:
        print(f"尝试从浏览器读取 cookies: {args.cookies_from_browser}（若失败请使用 --cookiefile 导出后再试）")

    with YoutubeDL(ydl_opts) as ydl:
        try:
            ydl.download(urls)
        except Exception as e:
            print("下载过程中出现异常：", e)
            print("提示：若错误与 YouTube 登录/反机器人相关，请尝试：")
            print("  1) 使用 CLI: yt-dlp --cookies-from-browser chrome -a youtube_urls.txt")
            print("  2) 或 用浏览器导出 cookies（Netscape 格式）并传入 --cookiefile /path/to/cookies.txt")
            raise

if __name__ == '__main__':
    main()
